Integrate DSA training into Megatron-LM - #1
Conversation
Integrates DSA into MultiLatentAttention and TransformerConfig. Co-authored-by: yuli <yuli@aisingapore.org>
|
Cursor Agent can help with this pull request. Just |
| def _hadamard_transform(x: Tensor) -> Tensor: | ||
| """Apply an in-place Hadamard transform along the last dimension.""" | ||
| hidden = x.size(-1) | ||
| if hidden & (hidden - 1) != 0: | ||
| raise ValueError( | ||
| f"DSA requires index head dimension to be a power of two, but got {hidden}." | ||
| ) | ||
| stride = 1 | ||
| out = x | ||
| while stride < hidden: | ||
| out = out.reshape(*out.shape[:-1], -1, 2, stride) | ||
| first = out[..., 0, :] | ||
| second = out[..., 1, :] | ||
| out = torch.cat((first + second, first - second), dim=-1) | ||
| stride <<= 1 | ||
| return out / math.sqrt(hidden) |
There was a problem hiding this comment.
I am not sure if this implementation is correct. I tested it using the following Python script:
import torch
import torch.nn.functional as F
import scipy
def reference_hadamard(x: torch.Tensor, scale: float = 1.0):
"""
See https://github.com/Dao-AILab/fast-hadamard-transform
"""
dim = x.size(-1)
scale = dim**-0.5
return F.linear(x, torch.Tensor(scipy.linalg.hadamard(dim))) * scale
def cursor_hadamard(x: torch.Tensor, scale: float = 1.0):
"""
Hadamard transform function written by Cursor.
"""
hidden = x.size(-1)
if hidden & (hidden - 1) != 0:
raise ValueError(
f"DSA requires index head dimension to be a power of two, but got {hidden}."
)
stride = 1
out = x
while stride < hidden:
out = out.reshape(*out.shape[:-1], -1, 2, stride)
first = out[..., 0, :]
second = out[..., 1, :]
out = torch.cat((first + second, first - second), dim=-1)
stride <<= 1
return out / math.sqrt(hidden)
def main():
x = torch.randn(2, 2, 4)
y1 = reference_hadamard(x)
print("Reference result for hadamard transform")
print(y1)
y2 = cursor_hadamard(x)
print("Cursor result for hadamard transform")
print(y2)
print("Hello from dsa-cursor!")
if __name__ == "__main__":
main()
The output is the following:
Reference result for hadamard transform
tensor([[[ 0.2654, 1.8008, -0.8284, 0.7358],
[ 2.0578, 0.3078, -0.5755, -0.1308]],
[[ 0.3302, -0.1754, 0.2045, 0.0398],
[ 0.3369, -2.5096, -0.6018, -0.5389]]])
Traceback (most recent call last):
File "/data/projects/71001002/nkarthik/dsa_cursor/test.py", line 48, in <module>
main()
File "/data/projects/71001002/nkarthik/dsa_cursor/test.py", line 40, in main
y2 = cursor_hadamard(x)
^^^^^^^^^^^^^^^^^^
File "/data/projects/71001002/nkarthik/dsa_cursor/test.py", line 27, in cursor_hadamard
out = out.reshape(*out.shape[:-1], -1, 2, stride)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: shape '[2, 2, 2, -1, 2, 2]' is invalid for input of size 16
The reference implementation taken from here works but cursor's does not. It could be that cursor's version is not designed to handle the specific input that I used.
Perhaps we can directly use the hadamard transform implementation provided here instead of trying to write it ourselves?
|
|
||
| weights = self.weight_proj(hidden_states) * (self.index_heads ** -0.5) | ||
|
|
||
| raw = torch.einsum("bqhd,bkd->bqhk", q_latent, k_latent) |
There was a problem hiding this comment.
In the DeepSeek code, the q_latent and k_latent are first scaled to fp8 before calculating the index scores. Is this being done somewhere else in this implementation?
I am also wondering whether it is ok to proceed with implementation of the Lightning Indexer in bf16 in the first instance and work on the low precision part later?
| batch = hidden_states.size(1) | ||
| dsa_hidden = hidden_states.permute(1, 0, 2).contiguous() | ||
| if self.config.q_lora_rank is not None: | ||
| qr_states = q_compressed.permute(1, 0, 2).contiguous() |
There was a problem hiding this comment.
The q_compressed does not seem to be defined anywhere in this function. It is actually computed in the get_query_key_value_tensors function later in this script. But that function does not return q_compressed. This part of the code might need to be restructured.
3 participants
What does this PR do ?
This PR introduces DeepSeek Sparse Attention (DSA) training capabilities to Megatron-LM, allowing models to learn and apply sparse attention masks during training.
Contribution process
flowchart LR A[Pre-checks] --> B[PR Tests] subgraph Code Review/Approval C1[Expert Review] --> C2[Final Review] end B --> C1 C2 --> D[Merge]Pre-checks
Core 0.8)Code review
The following process is enforced via the CODEOWNERS file for changes into
megatron/core. For changes outside ofmegatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.For MRs into `main` branch
(Step 1): Add PR label
Expert Review(Step 2): Collect the expert reviewers reviews
Expert Reviewlabel when your PR is ready for review.Final Review might get declined if these requirements are not fulfilled.
(Step 3): Final Review
Final Reviewlabel(Optional Step 4): Cherry-pick into release branch
If this PR also needs to be merged into
core_r*release branches, after this PR has been merged, selectCherry-pickto open a new PR into the release branch.For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.Merging your PR
Any member of core-adlr and
core-nemowill be able to merge your PR.